#include "codec/ffmpeg_decoder.h"
extern "C" {
#include "libavformat/avformat.h"
#include "libavcodec/avcodec.h"
#include "libavutil/avutil.h"
#include "libswscale/swscale.h"
#include "libavutil/avassert.h"
#include "libavutil/hwcontext.h"
#include "libavutil/imgutils.h"
#include "libavutil/opt.h"
#include <libavutil/pixdesc.h>
#include <libavutil/dict.h>

}
#include <stdexcept>


#include "perf.h"

struct FFmpegDecoder::Impl {
    AVFormatContext* fmt_ctx = nullptr;
    AVCodecContext* codec_ctx = nullptr;
    AVFrame* frame = nullptr;
    AVPacket* pkt = nullptr;
    SwsContext* sws_ctx = nullptr;
    int video_stream_idx = -1;
    int width = 0, height = 0;
    bool opened = false;
    enum AVHWDeviceType dev_type;
};


static enum AVPixelFormat hw_pix_fmt = AV_PIX_FMT_VSV;
static enum AVPixelFormat get_hw_format(AVCodecContext *ctx, const enum AVPixelFormat *pix_fmts) {
    const enum AVPixelFormat *p;
    for (p = pix_fmts; *p != -1; p++) {
        if (*p == hw_pix_fmt)
            return *p;
    }
    av_log(ctx, AV_LOG_ERROR, "Failed to get HW surface format.\n");
    return AV_PIX_FMT_NONE;
}


FFmpegDecoder::FFmpegDecoder() : impl_(new Impl) {
    av_register_all();
    avformat_network_init();
    av_log_set_level(AV_LOG_VERBOSE);
    
}

FFmpegDecoder::~FFmpegDecoder() {
    close();
    delete impl_;
}

bool FFmpegDecoder::open(const std::string& input_url) {
    close();
    auto& d = *impl_;
    const char *dev_type = "vsv";
#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(58, 18, 100)
    d.dev_type = av_hwdevice_find_type_by_name(dev_type);
    if (d.dev_type == AV_HWDEVICE_TYPE_NONE) {
        fprintf(stderr, "Device type %s is not supported.\n", dev_type);
        fprintf(stderr, "Available device types:");
        while ((d.dev_type = av_hwdevice_iterate_types(d.dev_type)) != AV_HWDEVICE_TYPE_NONE)
            fprintf(stderr, " %s", av_hwdevice_get_type_name(d.dev_type));
        fprintf(stderr, "\n");
        return -1;
    }
#endif

    AVCodec* decoder = NULL;
    if (avformat_open_input(&d.fmt_ctx, input_url.c_str(), nullptr, nullptr) < 0) return false;
    if (avformat_find_stream_info(d.fmt_ctx, nullptr) < 0) return false;
    for (unsigned i = 0; i < d.fmt_ctx->nb_streams; ++i) {
        if (d.fmt_ctx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
            d.video_stream_idx = i;
            break;
        }
    }
    if (d.video_stream_idx < 0) return false;
    AVCodecParameters* codecpar = d.fmt_ctx->streams[d.video_stream_idx]->codecpar;

    av_log(NULL, AV_LOG_DEBUG, "video codec_id is %d.\n", codecpar->codec_id);

    switch (codecpar->codec_id) {
        case AV_CODEC_ID_H264:
            decoder = avcodec_find_decoder_by_name("h264_vsv_decoder");
            break;
        case AV_CODEC_ID_HEVC:
            decoder = avcodec_find_decoder_by_name("hevc_vsv_decoder");
            break;
        case AV_CODEC_ID_VP9:
            decoder = avcodec_find_decoder_by_name("vp9_vsv_decoder");
            break;
        case AV_CODEC_ID_MJPEG:
            decoder = avcodec_find_decoder_by_name("jpeg_vsv_decoder");
            break;
        default:
            decoder = avcodec_find_decoder(codecpar->codec_id);
            break;
    }

    if (decoder == NULL) {
        fprintf(stderr, "Unsupported codec! \n");
        return -1;
    }

#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(58, 18, 100)
    for (int i = 0;; i++) {
        const AVCodecHWConfig* config = avcodec_get_hw_config(decoder, i);
        if (!config) {
            fprintf(stderr, "Decoder %s does not support device type %s.\n", decoder->name,
                    av_hwdevice_get_type_name(d.dev_type));
            return -1;
        }

        fprintf(stderr, "Decoder: %s, device type: %s.\n", decoder->name,
                av_hwdevice_get_type_name(d.dev_type));

        if (config->methods & AV_CODEC_HW_CONFIG_METHOD_HW_DEVICE_CTX && config->device_type == d.dev_type) {
            hw_pix_fmt = config->pix_fmt;
            fprintf(stderr, "hw_pix_fmt %d\n", hw_pix_fmt);
            break;
        }
    }
#endif
    if (!(d.codec_ctx = avcodec_alloc_context3(decoder))) return AVERROR(ENOMEM);
    if (avcodec_parameters_to_context(d.codec_ctx, codecpar) < 0) return -1;
    d.codec_ctx->get_format = get_hw_format;

    AVDictionary *dec_opts=NULL;
    AVDictionary *opts = NULL;

    int card_id = 0;
    int vid = 0;
    av_dict_set_int(&dec_opts, "card_id", card_id, 0);
    av_dict_set_int(&dec_opts, "vpu_id", vid, 0);
    av_dict_set(&dec_opts, "output_pixfmt", "yuv420p", 0);

    char tmp[128];
    sprintf(tmp, "/dev/gcu%dvid%d", card_id, vid);
    av_dict_set(&opts, "dec", tmp, 0);
    av_dict_set(&opts, "enc", tmp, 0);
    av_dict_set(&opts, "mem", "/dev/gcu0", 0);
    av_dict_set(&opts, "mapped_io", "1", 0);

    // QSV硬件解码初始化
    AVBufferRef* hw_device_ctx = nullptr;
    if (av_hwdevice_ctx_create(&hw_device_ctx, d.dev_type, nullptr, opts, 0) == 0) {
        d.codec_ctx->hw_device_ctx = av_buffer_ref(hw_device_ctx);
        d.codec_ctx->get_format = get_hw_format;
    }else{
        fprintf(stdout, "Failed to create specified HW device.\n");
        return -1;
    }

    d.codec_ctx->pix_fmt = AV_PIX_FMT_BGR24;
    //av_opt_set(d.codec_ctx->priv_data, "pp_set", "640x360:0:0:1920x1080", 0);

    if (avcodec_open2(d.codec_ctx, decoder, &dec_opts) < 0) return false;
    
    av_dict_free(&dec_opts);
    av_dict_free(&opts);

    d.frame = av_frame_alloc();
    d.pkt = av_packet_alloc();
    d.width = d.codec_ctx->width;
    d.height = d.codec_ctx->height;
    d.sws_ctx = sws_getContext(d.width, d.height, AV_PIX_FMT_YUV420P, d.width, d.height, AV_PIX_FMT_BGR24, SWS_BILINEAR, nullptr, nullptr, nullptr);
    d.opened = true;
    return true;
}

bool FFmpegDecoder::readFrame(cv::Mat& frame) {

    //Perf perf("FFmpegDecoder::readFrame");
    auto& d = *impl_;
    if (!d.opened) return false;
    while (true) {
        int read_ret = av_read_frame(d.fmt_ctx, d.pkt);
        if (read_ret < 0) {
            // 文件结束，自动seek到头部
            av_seek_frame(d.fmt_ctx, d.video_stream_idx, 0, AVSEEK_FLAG_BACKWARD);
            continue;
        }

        //
        if (d.pkt->stream_index == d.video_stream_idx) {
            //perf.begin("FFmpegDecoder::readFrame:send_packet");
            int ret = avcodec_send_packet(d.codec_ctx, d.pkt);
            //perf.end("FFmpegDecoder::readFrame:send_packet");
            av_packet_unref(d.pkt);
            if (ret < 0) continue;
            //perf.begin("FFmpegDecoder::readFrame:receive_loop");
            while(ret >=0) {
                //perf.begin("FFmpegDecoder::avcodec_receive_frame");
                ret = avcodec_receive_frame(d.codec_ctx, d.frame);
                //perf.end("FFmpegDecoder::avcodec_receive_frame");
                if (ret == AVERROR_EOF)
                {
                    //解码完成，推出循环
                    printf("ret AVERROR_EOF = %d\n", ret);
                    break;
                }
                else if (ret == AVERROR(EAGAIN))
                {
                    //需要更多数据
                    //printf("ret AVERROR(EAGAIN) = %d\n", ret);
                    break;
                }
                else if (ret < 0)
                {
                    //解码发生错误
                    printf("Error while decoding, ret=%d\n", ret);
                    break;
                }

                // if ret > 0, success
                AVFrame *sw_frame = av_frame_alloc();
                //printf("d.frame->format=%d, d.ctx.pixfmt=%d\n", d.frame->format, d.codec_ctx->pix_fmt);
                if (d.codec_ctx->pix_fmt == hw_pix_fmt)
                {
                    // QSV硬件帧转内存帧
                    printf("hardware D->H\n");
                    //perf.begin("FFmpegDecoder::av_hwframe_transfer_data");
                    av_hwframe_transfer_data(sw_frame, d.frame, 0);
                    //perf.end("FFmpegDecoder::av_hwframe_transfer_data");
                }
                else
                {
                    av_frame_ref(sw_frame, d.frame);
                }
#if 1
                cv::Mat img(sw_frame->height, sw_frame->width, CV_8UC3);
                uint8_t *dst[] = {img.data};
            
                int dst_linesize[] = {static_cast<int>(img.step)};
                //printf("@@@@@@@\n");
                //perf.begin("FFmpegDecoder::sws_scale");
                //sws_scale(d.sws_ctx, d.frame->data, d.frame->linesize, 0, d.height, dst, dst_linesize);
                ret = av_image_copy_to_buffer(dst[0], img.cols*img.rows*3,
                                       (const uint8_t * const *)sw_frame->data,
                                       (const int *)sw_frame->linesize, AV_PIX_FMT_BGR24,
                                       sw_frame->width, sw_frame->height, 1);
                frame = img;
                //perf.end("FFmpegDecoder::sws_scale");
                //printf("222222222222\n");
#endif
                av_frame_free(&sw_frame);
            }
            //perf.end("FFmpegDecoder::readFrame:receive_loop");
            return true;
        } else {
            av_packet_unref(d.pkt);
        }
    }
    return false;
}

void FFmpegDecoder::close() {
    auto& d = *impl_;
    if (d.frame) av_frame_free(&d.frame);
    if (d.pkt) av_packet_free(&d.pkt);
    if (d.codec_ctx) avcodec_free_context(&d.codec_ctx);
    if (d.fmt_ctx) avformat_close_input(&d.fmt_ctx);
    if (d.sws_ctx) sws_freeContext(d.sws_ctx);
    d.opened = false;
}

bool FFmpegDecoder::isOpened() const { return impl_->opened; }
